#include <stdio.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

#define CHECK_RC(rc, message, handle) \
    do { \
        if ((rc) != TSS2_RC_SUCCESS) { \
            fprintf(stderr, "错误: %s 失败! TSS2_RC: 0x%x, %s\n", \
                    (message), (rc), Tss2_RC_Decode(rc)); \
            goto cleanup; \
        } \
    } while (0)

typedef struct {
    TPM2_ALG_ID iid;
    const char *name;
} AlgMapEntry;

const AlgMapEntry alg_map[] = {
    {TPM2_ALG_RSA,              "RSA"           },
    {TPM2_ALG_TDES,             "TDES"          },
    {TPM2_ALG_SHA,              "SHA"           },
    {TPM2_ALG_SHA1,             "SHA1"          },
    {TPM2_ALG_HMAC,             "HMAC"          },
    {TPM2_ALG_AES,              "AES"           },
    {TPM2_ALG_MGF1,             "MGF1"          },
    {TPM2_ALG_KEYEDHASH,        "KeyedHash"     },
    {TPM2_ALG_XOR,              "XOR"           },
    {TPM2_ALG_SHA256,           "SHA256"        },
    {TPM2_ALG_SHA384,           "SHA384"        },
    {TPM2_ALG_SHA512,           "SHA512"        },
    {TPM2_ALG_NULL,             "NULL"          },
    {TPM2_ALG_SM3_256,          "SM3_256"       },
    {TPM2_ALG_SM4,              "SM4"           },
    {TPM2_ALG_RSASSA,           "RSASSA"        },
    {TPM2_ALG_RSAES,            "RSAES"         },
    {TPM2_ALG_RSAPSS,           "RSAPSS"        },
    {TPM2_ALG_OAEP,             "OAEP"          },
    {TPM2_ALG_ECDSA,            "ECDSA"         },
    {TPM2_ALG_ECDH,             "ECDH"          },
    {TPM2_ALG_ECDAA,            "ECDAA"         },
    {TPM2_ALG_SM2,              "SM2"           },
    {TPM2_ALG_ECSCHNORR,        "ECSchnorr"     },
    {TPM2_ALG_ECMQV,            "ECMQV"         },
    {TPM2_ALG_KDF1_SP800_56A,   "KDF1_SP800_56A"},
    {TPM2_ALG_KDF2,             "KDF2"          },
    {TPM2_ALG_KDF1_SP800_108,   "KDF1_SP800_108"},
    {TPM2_ALG_ECC,              "ECC"           },
    {TPM2_ALG_SYMCIPHER,        "SymCipher"     },
    {TPM2_ALG_CAMELLIA,         "Camellia"      },
    {TPM2_ALG_CMAC,             "CMAC"          },
    {TPM2_ALG_CTR,              "CTR"           },
    {TPM2_ALG_SHA3_256,         "SHA3_256"      },
    {TPM2_ALG_SHA3_384,         "SHA3_384"      },
    {TPM2_ALG_SHA3_512,         "SHA3_512"      },
    {TPM2_ALG_OFB,              "OFB"           },
    {TPM2_ALG_CBC,              "CBC"           },
    {TPM2_ALG_CFB,              "CFB"           },
    {TPM2_ALG_ECB,              "ECB"           },
};

const char* algid_to_name(TPM2_ALG_ID alg);
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context);
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT **sys_context);
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context);
void clean(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT *sys_context);
TSS2_RC get_cap_algs(TSS2_SYS_CONTEXT *sys_context);

int main() {
    TSS2_RC rc = TSS2_RC_SUCCESS;
    TSS2_TCTI_CONTEXT *tcti_context = NULL;
    TSS2_SYS_CONTEXT *sys_context = NULL;

    /* --- 初始化 TCTI Context --- */
    rc = init_tcti_context(&tcti_context);
    CHECK_RC(rc, "初始化 TCTI Context", tcti_context);

    /* --- 初始化 SAPI Context --- */
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_RC(rc, "初始化 SAPI Context", sys_context);

    /* --- 启动 TPM --- */
    rc = startup_tpm(sys_context);
    CHECK_RC(rc, "启动 TPM", sys_context);

    /* --- 查询 TPM 支持的算法 --- */
    rc = get_cap_algs(sys_context);
    CHECK_RC(rc, "查询 TPM 支持的算法", sys_context);

    /* --- 清理资源 --- */
cleanup:
    clean(tcti_context, sys_context);

    return !(rc == TSS2_RC_SUCCESS);
}

// 初始化 TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    return Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
}

// 初始化 SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size, tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// 启动 TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc = Tss2_Sys_GetCapability(sys_context, NULL,
                                        TPM2_CAP_TPM_PROPERTIES, TPM2_PT_MANUFACTURER,
                                        1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

const char* algid_to_name(TPM2_ALG_ID alg) {
    for (size_t i = 0; i < sizeof(alg_map) / sizeof(AlgMapEntry); i++) {
        if (alg_map[i].iid == alg) {
            return alg_map[i].name;
        }
    }
    return "Unknown Algorithm";
}

TSS2_RC get_cap_algs(TSS2_SYS_CONTEXT *sys_context) {
    TSS2_RC rc;
    TPMI_YES_NO more_data;
    TPMS_CAPABILITY_DATA cap_data;
    TPM2_ALG_ID next_alg_id = TPM2_ALG_FIRST;

    printf("TPM 支持的算法:\n");
    printf("  AlgID   AlgName\n");
    printf("  -----------------\n");
    do {
        rc = Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_ALGS,
                                    next_alg_id, TPM2_MAX_CAP_ALGS, 
                                    &more_data, &cap_data, NULL);

        if (rc != TSS2_RC_SUCCESS) { return rc; }
        if (cap_data.data.algorithms.count == 0) { break; }

        for (UINT32 i = 0; i < cap_data.data.algorithms.count; ++i) {
            TPM2_ALG_ID alg = cap_data.data.algorithms.algProperties[i].alg;
            printf("  0x%04X  %s\n", alg, algid_to_name(alg));
            next_alg_id = alg + 1;
        }

    } while (more_data == TPM2_YES); 

    return TSS2_RC_SUCCESS;
}

// 清理资源
void clean(TSS2_TCTI_CONTEXT *tcti_context, TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
    }
}
